import os
"""The operator 'aten::avg_pool3d.out' is not currently implemented for the MPS device."""
os.environ['PYTORCH_ENABLE_MPS_FALLBACK']='1'

import torch
from torch import nn


IMG_SIZE = 224
LR = 0.001


class AlexNet(nn.Module):
    """ AlexNet
    Paper: ImageNet Classification with Deep Convolutional Neural Networks
    URL: https://papers.nips.cc/paper_files/paper/2012/file/c399862d3b9d6b76c8436e924a68c45b-Paper.pdf
    网络架构：
        第一层：卷积层，使用 96 个 11×11 像素的卷积核，步长为 4 像素。
        第二层：ReLU激活函数。
        第三层：LRN, size=5, k = 2, n = 5, α = 1e−4, and β = 0.75
        第四层：最大池化层，使用 3×3 像素的窗口，步长为 2 像素。
        第五层：卷积层，使用 256 个 5×5 像素的卷积核。
        第六层：ReLU激活函数。
        第七层：LRN, size=5, k = 2, n = 5, α = 1e−4, and β = 0.75
        第八层：最大池化层，使用 3×3 像素的窗口，步长为 2 像素。
        第九层：卷积层，使用 384 个 3×3 像素的卷积核。
        第十层：卷积层，使用 384 个 3×3 像素的卷积核，但输入通道数翻倍。
        第十一层：卷积层，使用 256 个 3×3 像素的卷积核。
        第十二：ReLU激活函数。
        第十三层：最大池化层，使用 3×3 像素的窗口，步长为 2 像素。
        第十四层：全连接层，有 4096 个神经元。
        第十五层：ReLU激活函数。
        第十六层：全连接层，有 4096 个神经元。
        第十七层：Dropout。
        第十八层：全连接层，有 1000 个神经元，输出层使用 softmax 函数。
    """
    def __init__(self, num_classes, in_channels=3):
        super(AlexNet, self).__init__()
        self.net = nn.Sequential(
            nn.Conv2d(in_channels, 96, kernel_size=11, stride=4, padding=0), nn.ReLU(), 
            nn.LocalResponseNorm(size=5, alpha=0.0001, beta=0.75, k=2),
            nn.MaxPool2d(kernel_size=3, stride=2),
            nn.Conv2d(96, 256, kernel_size=5, stride=1, padding=2), nn.ReLU(),
            nn.LocalResponseNorm(size=5, alpha=0.0001, beta=0.75, k=2),
            nn.MaxPool2d(kernel_size=3, stride=2),
            nn.Conv2d(256, 384, kernel_size=5, stride=1, padding=1), nn.ReLU(),
            nn.Conv2d(384, 384, kernel_size=3, stride=1, padding=1), nn.ReLU(),
            nn.Conv2d(384, 256, kernel_size=3, stride=1, padding=1), nn.ReLU(),
            nn.MaxPool2d(kernel_size=3, stride=2),
            nn.Flatten(),
            nn.Linear(256 * 4 * 4, 4096), nn.ReLU(), nn.Dropout(),
            nn.Linear(4096, 4096), nn.ReLU(), nn.Dropout(),
            nn.Linear(4096, num_classes)
        )
    

    def forward(self, input):
        return self.net(input)
    

    @property
    def image_size(self):
        return IMG_SIZE
    

    @property
    def lr(self):
        return LR